import torch
from ultralytics import YOLO
import glob
import os


class Yolov8Detect():
    def __init__(self, weights):
        cuda = True if torch.cuda.is_available() else False
        self.device = torch.device('cuda:0' if cuda else 'cpu')
        self.detect_model = YOLO(weights)
        self.detect_model.to(self.device)

    def inferences(self, inputs):
        results = self.detect_model(inputs)
        for result in results:
            label_text = []
            boxes = result.boxes
            for box in boxes:
                cat_num = int(box.cls.cpu())
                label_text.append([cat_num, box.xywhn.cpu().numpy()])
            # 修正保存路径处理
            save_path = inputs.replace('.jpg', '.txt').replace('images', 'labels')
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            txt_construct(save_path, label_text=label_text)


def txt_construct(save_path, label_text):
    with open(save_path, 'w') as file:
        file.truncate()
    for label in label_text:
        with open(save_path, 'a') as txt_file:
            label_ = label[0]
            size = label[1][0].tolist()
            size_string = ' '.join(map(str, size))
            result = f'{label_} {size_string}'
            print('result', result)
            txt_file.write(str(result))
            txt_file.write('\n')


if __name__ == '__main__':
    model_path = r'E:\Users\kx15\Desktop\practise\hole\model\best.pt'
    model = Yolov8Detect(model_path)

    # 修正图像路径
    image_folder = r'J:\20250904万州采集靶纸\images'
    image_paths = glob.glob(os.path.join(image_folder, '*.jpg'))

    for img_path in image_paths:
        model.inferences(img_path)
